# -*- coding:utf-8 -*-
"""Contains a variant of the densenet model definition."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

import tensorflow.contrib.slim as slim
import tensorflow.contrib.layers as layers

layers.variance_scaling_initializer()


def trunc_normal(stddev): return tf.truncated_normal_initializer(stddev=stddev)


def bn_act_conv_drp(current, num_outputs, kernel_size, scope='block'):
    current = slim.batch_norm(current, scope=scope + '_bn')
    current = tf.nn.relu(current)
    current = slim.conv2d(current, num_outputs, kernel_size, scope=scope + '_conv')
    current = slim.dropout(current, scope=scope + '_dropout')
    return current


def block(net, layers, growth, scope='block'):
    for idx in range(layers):
        bottleneck = bn_act_conv_drp(net, 4 * growth, [1, 1],
                                     scope=scope + '_conv1x1' + str(idx))
        tmp = bn_act_conv_drp(bottleneck, growth, [3, 3],
                              scope=scope + '_conv3x3' + str(idx))
        net = tf.concat(axis=3, values=[net, tmp])
    return net


def transition(net, compression_ratio=0.5):
    output = int(net.shape[3].value * compression_ratio)
    # 1*1卷积压缩
    net = slim.conv2d(net, kernel_size=[1, 1], num_outputs=output)
    # 2*2平均池化降低feature map维度
    net = slim.avg_pool2d(net, kernel_size=[2, 2], stride=2, padding='VALID')

    return net


def densenet(images, num_classes=10, is_training=False,
             dropout_keep_prob=0.8, weight_decay=0.0001,
             scope='densenet'):
    """Creates a variant of the densenet model.

      images: A batch of `Tensors` of size [batch_size, height, width, channels].
      num_classes: the number of classes in the dataset.
      is_training: specifies whether or not we're currently training the model.
        This variable will determine the behaviour of the dropout layer.
      dropout_keep_prob: the percentage of activation values that are retained.
      prediction_fn: a function to get predictions out of logits.
      scope: Optional variable_scope.

    Returns:
      logits: the pre-softmax activations, a tensor of size
        [batch_size, `num_classes`]
      end_points: a dictionary from components of the network to the corresponding
        activation.
    """
    growth = 32
    compression_rate = 0.5
    layers_per_block = 16

    def reduce_dim(input_feature):
        return int(int(input_feature.shape[-1]) * compression_rate)

    end_points = {}

    with tf.variable_scope(scope, 'DenseNet', [images, num_classes]):
        with slim.arg_scope(densenet_arg_scope(is_training=is_training,
                                               keep_prob=dropout_keep_prob,
                                               weight_decay=0.0001)):
            pass
            # TODO: DenseNet实现
            '''
            采用DenseNet-BC模型 总层数250，增长率24，压缩率0.5
            采用三个Dense blocks每个blocks包含82层
            feature map shape分别是32*32， 16*16，8*8
            '''

            # 首先进行一个2倍增长率输出的卷积 卷积核为1*1
            net = slim.conv2d(inputs=images, num_outputs=2 * growth, kernel_size=[7, 7], stride=2)

            # 第一个dense block
            net = block(net=net, layers=6, growth=growth, scope='block1')

            # 第一个Transition Layer
            net = transition(net)

            # 第二个dense block
            net = block(net, 12, growth, 'block2')
            # 第二个Transition Layer
            net = transition(net)

            # 第三个dense block
            net = block(net, 24, growth, 'block3')
            # 第二个Transition Layer
            net = transition(net)

            # 第三个dense block
            net = block(net, 16, growth, 'block4')
            # 第二个Transition Layer
            net = transition(net)

            # Classification Layer 全局平均池化
            net = slim.avg_pool2d(net, [7, 7], padding='VALID')
            # 降维为2介张量
            # net = slim.flatten(net)
            net = tf.squeeze(net)
            # 全连接层生成logits
            logits = slim.fully_connected(net, 200)

            end_points['logits'] = logits

    return logits, end_points


# def bn_drp_scope(is_training=True, keep_prob=0.8):
#     keep_prob = keep_prob if is_training else 1
#     with slim.arg_scope(
#         [slim.batch_norm],
#             scale=True, is_training=is_training, updates_collections=None):
#         with slim.arg_scope(
#             [slim.dropout],
#                 is_training=is_training, keep_prob=keep_prob) as bsc:
#             return bsc


def densenet_arg_scope(weight_decay=0.0001, is_training=True, keep_prob=0.8):
    """"""
    """Defines the default densenet argument scope.

    Args:
      weight_decay: The weight decay to use for regularizing the model.
      is_training: Is training
      keep_prob: Keep prob

    Returns:
      An `arg_scope` to use for the densenet model.
    """
    keep_prob = keep_prob if is_training else 1
    with slim.arg_scope(
            [slim.batch_norm],
            scale=True, is_training=is_training, updates_collections=None):
        with slim.arg_scope(
                [slim.dropout],
                keep_prob=keep_prob, is_training=is_training):
            with slim.arg_scope(
                    [slim.conv2d, slim.fully_connected],
                    weights_initializer=tf.contrib.layers.variance_scaling_initializer(
                        factor=2.0, mode='FAN_IN', uniform=False),
                    weights_regularizer=slim.l2_regularizer(weight_decay),
                    activation_fn=None, biases_initializer=None):
                with slim.arg_scope(
                        [slim.conv2d],
                        padding='same',
                        stride=1
                ) as sc:
                    return sc


densenet.default_image_size = 224
